// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/spill.h>
#include <pollux/common/base/runtime_metrics.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/common/testutil/test_value.h>
#include <pollux/serializers/presto_serializer.h>

using kumo::pollux::common::testutil::TestValue;

namespace kumo::pollux::exec {
    namespace {
        // Returns the CompareFlags vector whose size is equal to numSortKeys. Fill in
        // with default CompareFlags() if 'compareFlags' is empty.
        const std::vector<CompareFlags> getCompareFlagsOrDefault(
            const std::vector<CompareFlags> &compareFlags,
            int32_t numSortKeys) {
            POLLUX_DCHECK(compareFlags.empty() || compareFlags.size() == numSortKeys);
            if (compareFlags.size() == numSortKeys) {
                return compareFlags;
            }
            return std::vector<CompareFlags>(numSortKeys);
        }
    } // namespace

    void SpillMergeStream::pop() {
        POLLUX_CHECK(!closed_);
        if (++index_ >= size_) {
            setNextBatch();
        }
    }

    int32_t SpillMergeStream::compare(const MergeStream &other) const {
        POLLUX_CHECK(!closed_);
        auto &otherStream = static_cast<const SpillMergeStream &>(other);
        auto &children = rowVector_->children();
        auto &otherChildren = otherStream.current().children();
        int32_t key = 0;
        if (sortCompareFlags().empty()) {
            do {
                auto result = children[key]
                        ->compare(
                            otherChildren[key].get(),
                            index_,
                            otherStream.index_,
                            CompareFlags())
                        .value();
                if (result != 0) {
                    return result;
                }
            } while (++key < numSortKeys());
        } else {
            do {
                auto result = children[key]
                        ->compare(
                            otherChildren[key].get(),
                            index_,
                            otherStream.index_,
                            sortCompareFlags()[key])
                        .value();
                if (result != 0) {
                    return result;
                }
            } while (++key < numSortKeys());
        }
        return 0;
    }

    void SpillMergeStream::close() {
        POLLUX_CHECK(!closed_);
        closed_ = true;
        rowVector_.reset();
        decoded_.clear();
        rows_.resize(0);
        index_ = 0;
        size_ = 0;
    }

    SpillState::SpillState(
        const common::GetSpillDirectoryPathCB &getSpillDirPathCb,
        const common::UpdateAndCheckSpillLimitCB &updateAndCheckSpillLimitCb,
        const std::string &fileNamePrefix,
        int32_t maxPartitions,
        int32_t numSortKeys,
        const std::vector<CompareFlags> &sortCompareFlags,
        uint64_t targetFileSize,
        uint64_t writeBufferSize,
        common::CompressionKind compressionKind,
        const std::optional<common::PrefixSortConfig> &prefixSortConfig,
        memory::MemoryPool *pool,
        melon::Synchronized<common::SpillStats> *stats,
        const std::string &fileCreateConfig)
        : getSpillDirPathCb_(getSpillDirPathCb),
          updateAndCheckSpillLimitCb_(updateAndCheckSpillLimitCb),
          fileNamePrefix_(fileNamePrefix),
          maxPartitions_(maxPartitions),
          numSortKeys_(numSortKeys),
          sortCompareFlags_(
              getCompareFlagsOrDefault(sortCompareFlags, numSortKeys)),
          targetFileSize_(targetFileSize),
          writeBufferSize_(writeBufferSize),
          compressionKind_(compressionKind),
          prefixSortConfig_(prefixSortConfig),
          fileCreateConfig_(fileCreateConfig),
          pool_(pool),
          stats_(stats),
          partitionWriters_(maxPartitions_) {
    }

    void SpillState::setPartitionSpilled(uint32_t partition) {
        POLLUX_DCHECK_LT(partition, maxPartitions_);
        POLLUX_DCHECK_LT(spilledPartitionSet_.size(), maxPartitions_);
        POLLUX_DCHECK(!spilledPartitionSet_.contains(partition));
        spilledPartitionSet_.insert(partition);
        ++stats_->wlock()->spilledPartitions;
        common::incrementGlobalSpilledPartitionStats();
    }

    /*static*/
    void SpillState::validateSpillBytesSize(uint64_t bytes) {
        static constexpr uint64_t kMaxSpillBytesPerWrite =
                std::numeric_limits<int32_t>::max();
        if (bytes >= kMaxSpillBytesPerWrite) {
            POLLUX_GENERIC_SPILL_FAILURE(fmt::format(
                "Spill bytes will overflow. Bytes {}, kMaxSpillBytesPerWrite: {}",
                bytes,
                kMaxSpillBytesPerWrite));
        }
    }

    void SpillState::updateSpilledInputBytes(uint64_t bytes) {
        auto statsLocked = stats_->wlock();
        statsLocked->spilledInputBytes += bytes;
        common::updateGlobalSpillMemoryBytes(bytes);
    }

    uint64_t SpillState::appendToPartition(
        uint32_t partition,
        const RowVectorPtr &rows) {
        POLLUX_CHECK(
            isPartitionSpilled(partition), "Partition {} is not spilled", partition);

        TestValue::adjust(
            "kumo::pollux::exec::SpillState::appendToPartition", this);

        POLLUX_CHECK_NOT_NULL(
            getSpillDirPathCb_, "Spill directory callback not specified.");
        auto spillDir = getSpillDirPathCb_();
        POLLUX_CHECK(!spillDir.empty(), "Spill directory does not exist");
        // Ensure that partition exist before writing.
        if (partitionWriters_.at(partition) == nullptr) {
            partitionWriters_[partition] = std::make_unique<SpillWriter>(
                std::static_pointer_cast<const RowType>(rows->type()),
                numSortKeys_,
                sortCompareFlags_,
                compressionKind_,
                fmt::format("{}/{}-spill-{}", spillDir, fileNamePrefix_, partition),
                targetFileSize_,
                writeBufferSize_,
                fileCreateConfig_,
                updateAndCheckSpillLimitCb_,
                pool_,
                stats_);
        }

        const uint64_t bytes = rows->estimate_flat_size();
        validateSpillBytesSize(bytes);
        updateSpilledInputBytes(bytes);

        IndexRange range{0, rows->size()};
        return partitionWriters_[partition]->write(
            rows, melon::Range<IndexRange *>(&range, 1));
    }

    SpillWriter *SpillState::partitionWriter(uint32_t partition) const {
        POLLUX_DCHECK(isPartitionSpilled(partition));
        return partitionWriters_[partition].get();
    }

    void SpillState::finishFile(uint32_t partition) {
        auto *writer = partitionWriter(partition);
        if (writer == nullptr) {
            return;
        }
        writer->finishFile();
    }

    size_t SpillState::numFinishedFiles(uint32_t partition) const {
        if (!isPartitionSpilled(partition)) {
            return 0;
        }
        const auto *writer = partitionWriter(partition);
        if (writer == nullptr) {
            return 0;
        }
        return writer->numFinishedFiles();
    }

    SpillFiles SpillState::finish(uint32_t partition) {
        auto *writer = partitionWriter(partition);
        if (writer == nullptr) {
            return {};
        }
        return writer->finish();
    }

    const SpillPartitionNumSet &SpillState::spilledPartitionSet() const {
        return spilledPartitionSet_;
    }

    std::vector<std::string> SpillState::testingSpilledFilePaths() const {
        std::vector<std::string> spilledFiles;
        for (const auto &writer: partitionWriters_) {
            if (writer != nullptr) {
                const auto partitionSpilledFiles = writer->testingSpilledFilePaths();
                spilledFiles.insert(
                    spilledFiles.end(),
                    partitionSpilledFiles.begin(),
                    partitionSpilledFiles.end());
            }
        }
        return spilledFiles;
    }

    std::vector<uint32_t> SpillState::testingSpilledFileIds(
        int32_t partitionNum) const {
        return partitionWriters_[partitionNum]->testingSpilledFileIds();
    }

    SpillPartitionNumSet SpillState::testingNonEmptySpilledPartitionSet() const {
        SpillPartitionNumSet partitionSet;
        for (uint32_t partition = 0; partition < maxPartitions_; ++partition) {
            if (partitionWriters_[partition] != nullptr) {
                partitionSet.insert(partition);
            }
        }
        return partitionSet;
    }

    std::vector<std::unique_ptr<SpillPartition> > SpillPartition::split(
        int numShards) {
        std::vector<std::unique_ptr<SpillPartition> > shards(numShards);
        const auto numFilesPerShard = files_.size() / numShards;
        int32_t numRemainingFiles = files_.size() % numShards;
        int fileIdx{0};
        for (int shard = 0; shard < numShards; ++shard) {
            SpillFiles files;
            auto numFiles = numFilesPerShard;
            if (numRemainingFiles-- > 0) {
                ++numFiles;
            }
            files.reserve(numFiles);
            while (files.size() < numFiles) {
                files.push_back(std::move(files_[fileIdx++]));
            }
            shards[shard] = std::make_unique<SpillPartition>(id_, std::move(files));
        }
        POLLUX_CHECK_EQ(fileIdx, files_.size());
        files_.clear();
        return shards;
    }

    std::string SpillPartition::toString() const {
        return fmt::format(
            "SPILLED PARTITION[ID:{} FILES:{} SIZE:{}]",
            id_.toString(),
            files_.size(),
            succinctBytes(size_));
    }

    std::unique_ptr<UnorderedStreamReader<BatchStream> >
    SpillPartition::createUnorderedReader(
        uint64_t bufferSize,
        memory::MemoryPool *pool,
        melon::Synchronized<common::SpillStats> *spillStats) {
        POLLUX_CHECK_NOT_NULL(pool);
        std::vector<std::unique_ptr<BatchStream> > streams;
        streams.reserve(files_.size());
        for (auto &fileInfo: files_) {
            streams.push_back(FileSpillBatchStream::create(
                SpillReadFile::create(fileInfo, bufferSize, pool, spillStats)));
        }
        files_.clear();
        return std::make_unique<UnorderedStreamReader<BatchStream> >(
            std::move(streams));
    }

    std::unique_ptr<TreeOfLosers<SpillMergeStream> >
    SpillPartition::createOrderedReader(
        uint64_t bufferSize,
        memory::MemoryPool *pool,
        melon::Synchronized<common::SpillStats> *spillStats) {
        std::vector<std::unique_ptr<SpillMergeStream> > streams;
        streams.reserve(files_.size());
        for (auto &fileInfo: files_) {
            streams.push_back(FileSpillMergeStream::create(
                SpillReadFile::create(fileInfo, bufferSize, pool, spillStats)));
        }
        files_.clear();
        // Check if the partition is empty or not.
        if (MELON_UNLIKELY(streams.empty())) {
            return nullptr;
        }
        return std::make_unique<TreeOfLosers<SpillMergeStream> >(std::move(streams));
    }

    uint32_t FileSpillMergeStream::id() const {
        POLLUX_CHECK(!closed_);
        return spillFile_->id();
    }

    void FileSpillMergeStream::nextBatch() {
        POLLUX_CHECK(!closed_);
        index_ = 0;
        if (!spillFile_->nextBatch(rowVector_)) {
            size_ = 0;
            close();
            return;
        }
        size_ = rowVector_->size();
    }

    void FileSpillMergeStream::close() {
        POLLUX_CHECK(!closed_);
        SpillMergeStream::close();
        spillFile_.reset();
    }

    SpillPartitionIdSet toSpillPartitionIdSet(
        const SpillPartitionSet &partitionSet) {
        SpillPartitionIdSet partitionIdSet;
        partitionIdSet.reserve(partitionSet.size());
        for (auto &partitionEntry: partitionSet) {
            partitionIdSet.insert(partitionEntry.first);
        }
        return partitionIdSet;
    }

    namespace {
        tsan_atomic<uint32_t> &maxSpillInjections() {
            static tsan_atomic<uint32_t> maxInjections{0};
            return maxInjections;
        }

        tsan_atomic<uint32_t> &testingSpillPct() {
            static tsan_atomic<uint32_t> spillPct{0};
            return spillPct;
        }

        std::string &testingSpillPoolRegExp() {
            static std::string spillPoolRegExp{".*"};
            return spillPoolRegExp;
        }
    } // namespace

    TestScopedSpillInjection::TestScopedSpillInjection(
        int32_t spillPct,
        const std::string &poolRegExp,
        uint32_t maxInjections) {
        POLLUX_CHECK_EQ(injectedSpillCount(), 0);
        testingSpillPct() = spillPct;
        testingSpillPoolRegExp() = poolRegExp;
        maxSpillInjections() = maxInjections;
        injectedSpillCount() = 0;
    }

    TestScopedSpillInjection::~TestScopedSpillInjection() {
        testingSpillPct() = 0;
        injectedSpillCount() = 0;
        testingSpillPoolRegExp() = ".*";
        maxSpillInjections() = 0;
    }

    tsan_atomic<uint32_t> &injectedSpillCount() {
        static tsan_atomic<uint32_t> injectedCount{0};
        return injectedCount;
    }

    bool testingTriggerSpill(const std::string &pool) {
        // Put cheap check first to reduce CPU consumption in release code.
        if (testingSpillPct() <= 0) {
            return false;
        }

        if (injectedSpillCount() >= maxSpillInjections()) {
            return false;
        }

        if (melon::Random::rand32() % 100 > testingSpillPct()) {
            return false;
        }

        if (!pool.empty() && !RE2::FullMatch(pool, testingSpillPoolRegExp())) {
            return false;
        }

        ++injectedSpillCount();
        return true;
    }

    void removeEmptyPartitions(SpillPartitionSet &partitionSet) {
        auto it = partitionSet.begin();
        while (it != partitionSet.end()) {
            if (it->second->numFiles() > 0) {
                ++it;
            } else {
                it = partitionSet.erase(it);
            }
        }
    }
} // namespace kumo::pollux::exec
